import importlib.resources as pkg_resources
import json_repair
from typing import List
from jinja2 import Environment
from pathlib import Path
from research_agent.core.survey import Survey
from research_agent.core.config import Config
from research_agent.core.general_llm import LLM
from pyaml_env import parse_config
import asyncio
import re
import logging
logger = logging.getLogger(__name__)

class PaperReviewer:
    def __init__(self,language="English"):

        self.language = language
        configs = parse_config(Config.YAML_CONFIG)
        self.llm = LLM(config=configs[Config.DEFAULT_MODEL])
        self.max_retries = 3
        # 获取当前文件所在目录的路径
        base_path = Path(__file__).parent / "prompts"

        # 修改文件路径的获取方式
        paper_review_prompt_file = base_path / "paper_review.jinja"
        with open(paper_review_prompt_file, "r",encoding="utf-8") as f:
            self.paper_review_prompt_template = Environment().from_string(f.read())

        paper_proofreading_prompt_file = base_path / "paper_proofreading.jinja"
        with open(paper_proofreading_prompt_file, "r", encoding="utf-8") as f:
            self.paper_proofreading_prompt_template = Environment().from_string(f.read())

        section_review_prompt_file = base_path / "section_review.jinja"
        with open(section_review_prompt_file, "r", encoding="utf-8") as f:
            self.section_review_prompt_template = Environment().from_string(f.read())

    def set_language(self,language):
        self.language = language
    def get_language(self):
        return self.language

    async def review_paper(self, topic: str, paper: str,survey:Survey) -> dict:
        # 将paper按照小节解析出来，并和survey对应上
        parsed_paper = self.parse_draft(paper)
        for section in parsed_paper:
            section_code = section["chapter_num"]
            if isinstance(section["chapter_content"], str):
                section_content = section["chapter_content"]
            elif isinstance(section["chapter_content"], List):
                section_content = '\n'.join([content_i for content_i in section["chapter_content"] if content_i != ""])
            else:
                raise ValueError("section[\"chapter_content\"] must be a string or a list of strings")

            survey.update_full_content(section_code,section_content)
            if section.get("subsection_content",None)!=None and len(section["subsection_content"])>0:
                for subsection in section["subsection_content"]:
                    subsection_code = subsection["chapter_num"]
                    subsection_content = '\n'.join([content_i for content_i in subsection["chapter_content"] if content_i != ""])
                    survey.update_full_content(subsection_code,subsection_content)

        # 将对每个小节进行修改的任务包装成task，开始执行
        root = survey.get_root()
        # 从根节点得到section节点
        section_nodes = survey.full_content.successors(root)
        section_nodes = list(section_nodes)

        if not isinstance(section_nodes, list) or not section_nodes:
            raise ValueError("section_nodes must be a non-empty list of strings")
        tasks = [self.proofreading_section(survey, node_section, topic) for
                 node_section in section_nodes]
        await asyncio.gather(*tasks, return_exceptions=True)
        # 修改完以后，开始根据内容提review问题，并返回这些问题
        # 因为也不确定改出来效果如何，这段代码先不用
        '''
        tasks = [self.review_section(survey, node_section, topic) for
                 node_section in section_nodes]
        await asyncio.gather(*tasks, return_exceptions=True)

        prompt_messages = self._prepare_prompts(topic, paper)
        response = await self.llm.completion(prompt_messages)
        return json_repair.loads(response)
        '''
        return survey

    async def proofreading_section(self,survey, node_section, topic):
        if not isinstance(node_section, int):
            raise ValueError("section node number must be a int number")
        section_code = survey.full_content.nodes[node_section]['code']
        section_title = survey.full_content.nodes[node_section]['title']
        section_description = survey.full_content.nodes[node_section]['description']
        section_content = survey.full_content.nodes[node_section]['content']
        # previous_content = previous_content + '\n' + "## " + section_code + " " + section_title
        logger.info(
            f"proofreading section {section_code}:{survey.full_content.nodes[node_section]['title']}……")
        subsection_nodes = survey.full_content.successors(node_section)
        subsection_nodes = list(subsection_nodes)
        # 先审校section的内容
        if section_content.strip()!= '':
            prompt_messages = self._prepare_proofreading_section_prompt(
                topic=topic,subsection_code=section_code, code_level='section', section_title =section_title,
                section_description = section_description,section_content=section_content
            )
            for attempt in range(self.max_retries):
                try:
                    response = await self.llm.completion(prompt_messages)
                    survey.update_full_content(section_code,json_repair.loads(response)["new_content"])
                    break
                except Exception as e:
                    logger.error(f"proofreading section {section_code}:{section_title} 失败: {e}")
                    if attempt < self.max_retries - 1:
                        await asyncio.sleep(0.5)
        if len(subsection_nodes) > 0:
            # 处理有子章节的情况，即先审校子章节，再审校section的内容
            if not isinstance(subsection_nodes, list) or not subsection_nodes:
                raise ValueError("subsection_nodes must be a non-empty list of strings")
            tasks = [self.proofreading_subsection(survey,node_subsection,topic) for node_subsection in subsection_nodes]
            await asyncio.gather(*tasks, return_exceptions=True)
        return survey

    async def proofreading_subsection(self, survey, node_subsection, topic):
        if not isinstance(node_subsection, int):
            raise ValueError("section node number must be a int number")
        subsection_code = survey.full_content.nodes[node_subsection]['code']
        subsection_title = survey.full_content.nodes[node_subsection]['title']
        subsection_description = survey.full_content.nodes[node_subsection]['description']
        subsection_content = survey.full_content.nodes[node_subsection]['content']
        logger.info(
            f"proofreading subsection {subsection_code}:{survey.full_content.nodes[node_subsection]['title']}……")
        # 先审校section的内容
        prompt_messages = self._prepare_proofreading_section_prompt(
            topic=topic, subsection_code=subsection_code, code_level='subsection', section_title=subsection_title,
            section_description=subsection_description, section_content=subsection_content
        )

        for attempt in range(self.max_retries):
            try:
                response = await self.llm.completion(prompt_messages)
                survey.update_full_content(subsection_code, json_repair.loads(response)["new_content"])
                break
            except Exception as e:
                logger.error(f"proofreading subsection {subsection_code}:{subsection_title} 失败: {e}")
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(0.5)
        return survey
    async def review_section(self,survey, node_section, topic):
        if not isinstance(node_section, int):
            raise ValueError("section node number must be a int number")
        section_code = survey.full_content.nodes[node_section]['code']
        section_title = survey.full_content.nodes[node_section]['title']
        section_description = survey.full_content.nodes[node_section]['description']
        section_content = survey.full_content.nodes[node_section]['content']
        # previous_content = previous_content + '\n' + "## " + section_code + " " + section_title
        logger.info(
            f"review section {section_code}:{survey.full_content.nodes[node_section]['title']}……")
        subsection_nodes = survey.full_content.successors(node_section)
        subsection_nodes = list(subsection_nodes)
        # 先审校section的内容
        prompt_messages = self._prepare_proofreading_section_prompt(
            topic=topic,subsection_code=section_code, code_level='section', section_title =section_title,
            section_description = section_description,section_content=section_content
        )
        for attempt in range(self.max_retries):
            try:
                response = await self.llm.completion(prompt_messages)
                survey.update_full_content(section_code,json_repair.loads(response)["new_content"])
                break
            except Exception as e:
                logger.error(f"review section {section_code}:{section_title} 失败: {e}")
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(0.5)
        if len(subsection_nodes) > 0:
            # 处理有子章节的情况，即先审校子章节，再审校section的内容
            if not isinstance(subsection_nodes, list) or not subsection_nodes:
                raise ValueError("subsection_nodes must be a non-empty list of strings")
            tasks = [self.proofreading_subsection(survey,node_subsection,topic) for node_subsection in subsection_nodes]
            await asyncio.gather(*tasks, return_exceptions=True)
        return survey

    async def review_subsection(self, survey, node_subsection, topic):
        if not isinstance(node_subsection, int):
            raise ValueError("section node number must be a int number")
        subsection_code = survey.full_content.nodes[node_subsection]['code']
        subsection_title = survey.full_content.nodes[node_subsection]['title']
        subsection_description = survey.full_content.nodes[node_subsection]['description']
        subsection_content = survey.full_content.nodes[node_subsection]['content']
        logger.info(
            f"review subsection {subsection_code}:{survey.full_content.nodes[node_subsection]['title']}……")
        # 先审校section的内容
        prompt_messages = self._prepare_proofreading_section_prompt(
            topic=topic, subsection_code=subsection_code, code_level='subsection', section_title=subsection_title,
            section_description=subsection_description, section_content=subsection_content
        )
        for attempt in range(self.max_retries):
            try:
                response = await self.llm.completion(prompt_messages)
                survey.update_full_content(subsection_code, json_repair.loads(response)["new_content"])
                break
            except Exception as e:
                logger.error(f"review subsection {subsection_code}:{subsection_title} 失败: {e}")
                if attempt < self.max_retries - 1:
                    await asyncio.sleep(0.5)
        return survey

    @staticmethod
    def parse_draft(content:str):
        sections = []
        current_section = None
        content_buffer = []
        subsection_buffer = None

        # 匹配带编号的标题（兼容不同数量的#标记）
        title_re = re.compile(r'^#+\s+((\d+)(\.\d+)*)\s+(.*)$')

        for line in content.split('\n'):
            match = title_re.match(line)
            if match:
                chapter_num = match.group(1)
                main_num = match.group(2)
                sub_num = match.group(3)
                title_text = match.group(4)
                # 判断标题层级
                if '.' in chapter_num:  # 二级标题
                    if current_section and current_section['chapter_num'] == main_num:
                        # 保存前一个子章节内容
                        if subsection_buffer:
                            current_section['subsection_content'].append(subsection_buffer)
                        # 开始新子章节
                        subsection_buffer = {
                            'chapter_num': chapter_num,
                            'chapter_content': []
                        }
                    else:
                        # 处理异常情况：二级标题出现在一级之前
                        continue
                else:  # 一级标题
                    # 保存前一个主章节
                    if current_section:
                        # 处理未保存的子章节
                        if subsection_buffer:
                            current_section['subsection_content'].append(subsection_buffer)
                        # 合并主章节内容
                        current_section['chapter_content'] = '\n'.join(content_buffer).strip()
                        sections.append(current_section)
                    # 初始化新章节
                    current_section = {
                        'chapter_num': chapter_num,
                        'chapter_content': [],
                        'subsection_content': []
                    }
                    content_buffer = []
                    subsection_buffer = None
                continue  # 标题行不存入内容

            # 内容处理逻辑
            if current_section:
                if subsection_buffer:  # 二级标题内容
                    subsection_buffer['chapter_content'].append(line)
                else:  # 一级标题内容
                    content_buffer.append(line)

        # 处理最后一个章节
        if current_section:
            # 合并最终内容
            current_section['chapter_content'] = '\n'.join(content_buffer).strip()
            if subsection_buffer:
                current_section['subsection_content'].append({
                    'chapter_num': subsection_buffer['chapter_num'],
                    'chapter_content': '\n'.join(subsection_buffer['chapter_content']).strip()
                })
            # 清理空子章节
            if not current_section['subsection_content']:
                del current_section['subsection_content']
            sections.append(current_section)
        return sections
    def _prepare_proofreading_section_prompt(self,
            topic,subsection_code, code_level, section_title,
            section_description, section_content
        ):
        system_prompt = self.paper_proofreading_prompt_template.render(role="system",code_level=code_level)
        user_prompt = self.paper_proofreading_prompt_template.render(
            role="user",
            topic=topic,
            subsection_code=subsection_code,
            code_level=code_level,
            section_title=section_title,
            section_description=section_description,
            section_content=section_content
        )
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    def _prepare_review_section_prompt(self,
            topic,subsection_code, code_level, section_title,
            section_description, section_content
        ):
        system_prompt = self.section_review_prompt_template.render(role="system",code_level=code_level)
        user_prompt = self.section_review_prompt_template.render(
            role="user",
            topic=topic,
            subsection_code=subsection_code,
            code_level=code_level,
            section_title=section_title,
            section_description=section_description,
            section_content=section_content
        )
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    def _prepare_prompts(self, topic: str, paper: str) -> list[dict]:
        system_prompt = self.paper_review_prompt_template.render(role="system")
        user_prompt = self.paper_review_prompt_template.render(role="user", topic=topic, paper=paper)
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]

if __name__ == "__main__":
    paper_reviewer = PaperReviewer()
    with open("../../draft_iteration_output/20250227_152130/0.md","r",encoding="utf-8") as f:
        paper = f.read()
    print(PaperReviewer().parse_article(paper))